import pandas as pd
from typing import List

from constants import LINE_NUMBER


class PeriodPivotHelper:
    def __init__(self, default_source_data: str):
        self.default_source_data = default_source_data

    def build(
        self,
        sheet_name: str,
        base_cols: List[str],
        merge_keys: List[str],
        period_column: str,
        screening_value: str,
        outgroup_value: str,
        screening_prefix: str,
        outgroup_prefix: str,
        expand_after_line_number: bool = True,
    ) -> pd.DataFrame:
        df = pd.read_excel(self.default_source_data, sheet_name=sheet_name, dtype=str, keep_default_na=False, na_values=[])
        if df is None or df.empty:
            return pd.DataFrame()
        # 规范化列名并去重
        df.columns = [str(c).strip() for c in df.columns]
        df = df.loc[:, ~df.columns.duplicated()]

        # 若期别列不存在，直接跳过该模块的透视
        if period_column not in df.columns:
            return pd.DataFrame()

        # 仅保留需要的期别
        df = df[df[period_column].isin([screening_value, outgroup_value])].copy()
        if df.empty:
            return pd.DataFrame()

        # 计算扩展列（行号之后的所有列）
        if expand_after_line_number and LINE_NUMBER in df.columns:
            row_idx = df.columns.get_loc(LINE_NUMBER)
            period_cols: List[str] = [c for c in df.columns[row_idx + 1:].tolist() if c not in merge_keys and c != period_column]
        else:
            # 回退为除连接键与期别列外的所有列
            period_cols = [c for c in df.columns if c not in merge_keys and c != period_column]
        if not period_cols:
            return pd.DataFrame()

        base_available = [c for c in base_cols if c in df.columns]
        base_data = df[base_available].astype(str).drop_duplicates(subset=merge_keys, keep='last')

        # 出组记录用于限定需要输出的受试者
        outgroup_records = df[df[period_column] == outgroup_value][merge_keys].drop_duplicates()
        if outgroup_records.empty:
            return pd.DataFrame()

        # 构造筛选期与出组期数据集
        screening_data = df[df[period_column] == screening_value].copy().astype(str)
        screening_data = screening_data.merge(outgroup_records, on=merge_keys, how='inner')[merge_keys + period_cols]
        screening_data.columns = merge_keys + [f"{screening_prefix}_{col}" for col in period_cols]

        outgroup_data = df[df[period_column] == outgroup_value][merge_keys + period_cols].copy().astype(str)
        outgroup_data.columns = merge_keys + [f"{outgroup_prefix}_{col}" for col in period_cols]

        # 汇总结果，保证筛选期在前、出组期在后
        result = base_data.merge(outgroup_records, on=merge_keys, how='inner')
        result = result.merge(screening_data, on=merge_keys, how='left')
        result = result.merge(outgroup_data, on=merge_keys, how='left')

        # 按列顺序组织输出
        desired_cols = base_available + [f"{screening_prefix}_{c}" for c in period_cols] + [f"{outgroup_prefix}_{c}" for c in period_cols]
        return result[[c for c in desired_cols if c in result.columns]]